using UnityEngine;
using Unity.MLAgents;
using Unity.MLAgents.Actuators;
using Unity.MLAgentsExamples;
using Unity.MLAgents.Sensors;
using Random = UnityEngine.Random;
using System.Collections.Generic;
using System.IO;
using System.Diagnostics;
using System.Linq;
using UnityEditor;

public class H1mimicAgent : Agent
{
    public bool train = false;
    public bool replay = false;

    float[] uff = new float[19] { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ,0, 0, 0, 0, 0, 0, 0, 0, 0};
    float[] u = new float[19] { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ,0, 0, 0, 0, 0, 0, 0, 0, 0};
    float[] utotal = new float[19] { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ,0, 0, 0, 0, 0, 0, 0, 0, 0};
    //int[] trainid=new int[13]{3,9, 5,13,6,14,15,   3,9, 5,13,16,17};
    int[] trainid=new int[13]{5,5,13,13,   6,14,15,   5,5,13,13,   16,17};
    
    private List<string> actionFolders = new List<string>();
    public int motion_idx = 0;
    public string motion_name;
    int idx = 0;

    private string dofFilePath;
    private string rotFilePath;
    private string posFilePath;

    private List<float[]> allDofData = new List<float[]>();
    private List<float[]> allPosData = new List<float[]>();
    private List<float[]> allRotData = new List<float[]>();
    
    private int currentFrame = 0;  
    private bool isEndEpisode = false;
    
    Transform body;

    List<float> P0 = new List<float>();
    List<float> W0 = new List<float>();
    List<Transform> bodypart = new List<Transform>();
    Vector3 pos0;
    Quaternion rot0;
    Quaternion newRotation;
    Vector3 newPosition;
    ArticulationBody[] jh = new ArticulationBody[19];
    ArticulationBody[] arts = new ArticulationBody[40];
    ArticulationBody art0;
    int tt = 0;

    private bool _isClone = false; 
    void Start()
    {
        Time.fixedDeltaTime = 0.02f;

        /*SerializedObject tagManager = new SerializedObject(AssetDatabase.LoadAllAssetsAtPath("ProjectSettings/TagManager.asset"));
        SerializedProperty layers = tagManager.FindProperty("layers");
        SerializedProperty layer = layers.GetArrayElementAtIndex(15);
        int targetLayer = LayerMask.NameToLayer("robot");
        layer.stringValue = "robot";
        tagManager.ApplyModifiedProperties();
        Physics.IgnoreLayerCollision(15, 15, true);
        ChangeLayerRecursively(gameObject, 15);*/

        if (train && !_isClone) 
        {
            for (int i = 1; i < 24; i++)
            {
                GameObject clone = Instantiate(gameObject); 
                clone.transform.position = transform.position + new Vector3(i * 2f, 0, 0);
                clone.name = $"{name}_Clone_{i}"; 
                clone.GetComponent<H1mimicAgent>()._isClone = true; 
            }
        }
    }

    void ChangeLayerRecursively(GameObject obj, int targetLayer)
    {
        obj.layer = targetLayer;
        foreach (Transform child in obj.transform)ChangeLayerRecursively(child.gameObject, targetLayer);
    }

    public override void Initialize()
    {

        arts = this.GetComponentsInChildren<ArticulationBody>();
        int ActionNum = 0;
        for (int k = 0; k < arts.Length; k++)
        {
            if(arts[k].jointType.ToString() == "RevoluteJoint")
            {
                jh[ActionNum] = arts[k];
                ActionNum++;
            }
        }
        body = arts[0].GetComponent<Transform>();
        art0 = body.GetComponent<ArticulationBody>();

        pos0 = body.position;
        rot0 = body.rotation;
        art0.GetJointPositions(P0);
        art0.GetJointVelocities(W0);
                
        string streamingAssetsPath = Path.Combine(Application.streamingAssetsPath, "h1_dataset");
        LoadActionFolders(streamingAssetsPath);
        print(actionFolders[1]);
        LoadDataForAction(actionFolders[1]);

    }
    void LoadActionFolders(string basePath)
    {
        string[] folders = Directory.GetDirectories(basePath);
        actionFolders = folders.ToList();  
    }
    void LoadDataForAction(string actionFolder)
    {
        dofFilePath = Path.Combine(actionFolder, "dof.csv");
        rotFilePath = Path.Combine(actionFolder, "root_rot.csv");
        posFilePath = Path.Combine(actionFolder, "root_trans_offset.csv");
        allDofData = LoadDataFromFile(dofFilePath);
        allRotData = LoadDataFromFile(rotFilePath);
        allPosData = LoadDataFromFile(posFilePath);
    }

    List<float[]> LoadDataFromFile(string filePath)
    {
        List<float[]> dataList = new List<float[]>();
        try
        {
            string[] lines = File.ReadAllLines(filePath);
            foreach (string line in lines)
            {
                string[] values = line.Split(',');
                List<float> frameData = new List<float>();
                foreach (string value in values)
                {
                    if (float.TryParse(value.Trim(), out float parsedValue))frameData.Add(parsedValue);
                }
                dataList.Add(frameData.ToArray());
            }
        }
        catch (System.Exception e)
        {
            print("Error loading data from file " + filePath + ": " + e.Message);
        }

        return dataList;
    }
    public override void OnEpisodeBegin()
    {
        art0.TeleportRoot(pos0, rot0);
        art0.velocity = Vector3.zero;
        art0.angularVelocity = Vector3.zero;
        art0.SetJointPositions(P0);
        art0.SetJointVelocities(W0);
        for (int i = 0; i < 19; i++) u[i] = 0;
        for (int i = 0; i < 19; i++) uff[i] = 0;
        motion_idx=trainid[idx];
        LoadDataForAction(actionFolders[motion_idx]);
        motion_name = actionFolders[motion_idx].Replace("./Assets/Imitation/H1/dataset\\", "");
        currentFrame = 0;
        if(!train)
        {
            idx++;
            if(idx==13)idx=0;
        }
        if(train)idx=(int)Mathf.Floor(((Time.time%3900)/300));//train for 10 million steps first
        //if(train)idx=(int)Mathf.Floor(((Time.time%390)/30));//train for 2 million steps after

        tt=0;
        float[] currentDof = allDofData[1];
        float[] currentPos = allPosData[1];
        float[] currentRot = allRotData[1];
        Vector3 newPosition = new Vector3(-currentPos[1], currentPos[2], currentPos[0]);
        Quaternion newRotation = new Quaternion(
            currentRot[1], 
            -currentRot[2], 
            currentRot[0], 
            currentRot[3]
        );
        newPosition.x+=pos0.x;
        newPosition.z+=pos0.z;
        
        art0.TeleportRoot(newPosition, newRotation);
        art0.velocity = Vector3.zero;
        art0.angularVelocity = Vector3.zero;
        art0.immovable = true;
        for (int i = 0; i < 19; i++)
        {
            uff[i] = currentDof[i]* 180f / 3.14f;
            SetJointTargetDeg(jh[i], uff[i]);
        }
    }

    public override void CollectObservations(VectorSensor sensor)
    {
        sensor.AddObservation(EulerTrans(body.eulerAngles[0]) * 3.14f / 180f );//rad
        sensor.AddObservation(EulerTrans(body.eulerAngles[2]) * 3.14f / 180f );//rad
        sensor.AddObservation(body.InverseTransformDirection(art0.velocity));
        sensor.AddObservation(body.InverseTransformDirection(art0.angularVelocity));
        for (int i = 0; i < 19; i++)
        {
            sensor.AddObservation((jh[i].jointPosition[0]));
            sensor.AddObservation(jh[i].jointVelocity[0]);
        }
    }
    
    float EulerTrans(float eulerAngle)
    {
        if (eulerAngle <= 180)
            return eulerAngle;
        else
            return eulerAngle - 360f;
    }
    
    public override void OnActionReceived(ActionBuffers actionBuffers)
    {
        var continuousActions = actionBuffers.ContinuousActions;
        var kk = 0.9f;
        
        float kb = 30;
        if(replay)kb = 0;
        for (int i = 0; i < 19; i++)
        {
            u[i] = u[i] * kk + (1 - kk) * continuousActions[i];
            utotal[i] = kb * u[i] + uff[i];
            SetJointTargetDeg(jh[i], utotal[i]);
        }
       
    }

    void FixedUpdate()
    {
	///////////////feedforward///////////////////////////////////////////////////////////
	if (allDofData.Count > 0)
	{
		float[] currentDof = allDofData[currentFrame];
		float[] currentPos = allPosData[currentFrame];
        float[] currentRot = allRotData[currentFrame];
		for (int i = 0; i < 19; i++)uff[i] = currentDof[i]* 180f / 3.14f;

        Quaternion gymQuat = new Quaternion(
            currentRot[0], 
            currentRot[1], 
            currentRot[2], 
            currentRot[3]
        );
        Quaternion conversionQ = new Quaternion(0.5f, -0.5f, -0.5f, 0.5f);
        newRotation = conversionQ * gymQuat * Quaternion.Inverse(conversionQ);
		newPosition = new Vector3(-currentPos[1], currentPos[2], currentPos[0]);
        newRotation = new Quaternion(
            currentRot[1], 
            -currentRot[2], 
            currentRot[0], 
            currentRot[3]
        );

        newPosition.x+=pos0.x;
        newPosition.z+=pos0.z;

    	if(replay)art0.TeleportRoot(newPosition, newRotation);//***************************************************************************************************************
   		currentFrame = (currentFrame + 1) % allDofData.Count;
		
		if (currentFrame == 0)  // 当前帧回到了第一帧
		{
            EndEpisode();  // 结束回合
		    currentFrame = 0;  // 重置为第一帧
		}
		    
		if (isEndEpisode)
		{
	            currentFrame = 0;
	            isEndEpisode = false; 
		}
	}
	
	/////////////////rewards/////////////////////////////////////////////////////////
        tt++;
        var vel = body.InverseTransformDirection(art0.velocity);
        var wel = body.InverseTransformDirection(art0.angularVelocity);
        
        var live_reward = 1f;
        float rot_reward = 0;
        float pos_reward = 0;
        if(tt>3)
        {
            art0.immovable = false;
            rot_reward = - 0.01f * Quaternion.Angle(body.rotation, newRotation);
            pos_reward = - 1f * (body.position - newPosition).magnitude;
            if (Quaternion.Angle(body.rotation, newRotation)>30f || (body.position - newPosition).magnitude>0.3f)EndEpisode();
        }  
        var reward = live_reward + (rot_reward + pos_reward)*1;
        AddReward(reward);
    
    }

    void SetJointTargetDeg(ArticulationBody joint, float x)
    {
        var drive = joint.xDrive;
        drive.stiffness = 2000f;//180f;
        drive.damping = 200f;//8f;
        drive.forceLimit = 300f;//250f;// 33.5f;// 50f;
        drive.target = x;
        joint.xDrive = drive;
    }

    public override void Heuristic(in ActionBuffers actionsOut)
    {
        
    }
    
}



